import pymysql
import logging
from datetime import datetime

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class DatabaseManager:
    def __init__(self, host='122.51.66.115', port=3306, user='abcde', password='webai666', database='WebAI'):
        """初始化数据库连接"""
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.database = database
        self.connection = None
        
    def connect(self):
        """建立数据库连接"""
        try:
            if self.connection is None or not self.connection.open:
                self.connection = pymysql.connect(
                    host=self.host,
                    port=self.port,
                    user=self.user,
                    password=self.password,
                    database=self.database,
                    charset='utf8mb4',
                    cursorclass=pymysql.cursors.DictCursor
                )
            return self.connection
        except Exception as e:
            logger.error(f"数据库连接失败: {str(e)}")
            raise
    
    def disconnect(self):
        """关闭数据库连接"""
        if self.connection and self.connection.open:
            self.connection.close()
            self.connection = None
    
    def insert_generation_record(self, user_id, generation_method, prompt=None, image_url=None, task_id=None, request_id=None, status='RUN'):
        """
        插入生成记录到数据库
        
        :param user_id: 用户ID
        :param generation_method: 生成方式 (prompt/image)
        :param prompt: 提示词内容（当生成方式为prompt时有值）
        :param image_url: 图片链接（当生成方式为image时有值）
        :param task_id: 任务ID
        :param request_id: 请求ID
        :param status: 状态（默认为RUN）
        :return: 插入成功的记录ID
        """
        try:
            # 建立连接
            conn = self.connect()
            
            with conn.cursor() as cursor:
                # 准备SQL语句 - 符合最新表结构
                sql = """
                INSERT INTO generation_records 
                (user_id, generation_method, prompt, image_url, task_id, request_id, status, created_time, updated_time, is_display)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                """
                
                now = datetime.now()
                # 执行插入
                cursor.execute(sql, (
                    user_id,
                    generation_method,
                    prompt,
                    image_url,
                    task_id,
                    request_id,
                    status,
                    now,  # created_time
                    now,  # updated_time
                    1  # is_display默认为true
                ))
            
            # 提交事务
            conn.commit()
            
            logger.info(f"成功插入生成记录，用户ID: {user_id}, 任务ID: {task_id}")
            return cursor.lastrowid
            
        except Exception as e:
            logger.error(f"插入生成记录失败: {str(e)}")
            if conn and not conn.closed:
                conn.rollback()
            raise
        finally:
            self.disconnect()
    
    def get_generation_record(self, record_id):
        """获取生成记录"""
        try:
            conn = self.connect()
            
            with conn.cursor() as cursor:
                sql = "SELECT * FROM generation_records WHERE id = %s"
                cursor.execute(sql, (record_id,))
                result = cursor.fetchone()
            
            return result
            
        except Exception as e:
            logger.error(f"获取生成记录失败: {str(e)}")
            raise
        finally:
            self.disconnect()
    
    def update_generation_status(self, task_id, status):
        """更新生成记录状态"""
        try:
            conn = self.connect()
            
            with conn.cursor() as cursor:
                sql = "UPDATE generation_records SET status = %s WHERE task_id = %s"
                cursor.execute(sql, (status, task_id))
            
            conn.commit()
            logger.info(f"成功更新生成记录状态，任务ID: {task_id}, 新状态: {status}")
            return cursor.rowcount
            
        except Exception as e:
            logger.error(f"更新生成记录状态失败: {str(e)}")
            if conn and not conn.closed:
                conn.rollback()
            raise
        finally:
            self.disconnect()
    
    def insert_evaluation_report(self, model_3d_url, user_id, title, markdown_report, screenshot=None):
        """
        插入3D模型评估报告到数据库
        
        :param model_3d_url: 3D模型链接
        :param user_id: 用户ID
        :param title: 报告标题
        :param markdown_report: 模型生成的Markdown格式报告
        :param screenshot: 3D模型的截图URL
        :return: 插入成功的记录ID
        """
        try:
            # 建立连接
            conn = self.connect()
            
            with conn.cursor() as cursor:
                # 准备SQL语句，严格按照数据库表结构
                sql = """
                INSERT INTO 3d_model_evaluation_reports 
                (model_3d_url, user_id, title, markdown_report, screenshot, read_status, create_time)
                VALUES (%s, %s, %s, %s, %s, %s, %s)
                """
                
                now = datetime.now()
                # 执行插入
                cursor.execute(sql, (
                    model_3d_url,
                    user_id,
                    title,
                    markdown_report,
                    screenshot,
                    0,  # read_status默认为0（未读）
                    now  # create_time
                ))
            
            # 提交事务
            conn.commit()
            
            logger.info(f"成功插入评估报告，用户ID: {user_id}, 模型URL: {model_3d_url}")
            return cursor.lastrowid
            
        except Exception as e:
            logger.error(f"插入评估报告失败: {str(e)}")
            if conn and not conn.closed:
                conn.rollback()
            raise
        finally:
            self.disconnect()
